package com.example.scribe.backend;

import java.util.ArrayList;
import java.util.List;

import com.example.scribe.SyntaxNode;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;
import net.bytebuddy.jar.asm.MethodVisitor;
import net.bytebuddy.jar.asm.ClassWriter;
import net.bytebuddy.jar.asm.Opcodes;

final class FunctionCompiler {
    private final String functionName;
    private final List<String> parameterNames;
    private final SyntaxNode bodyNode;
    private final String generatedClassName;

    FunctionCompiler(String functionName, List<String> parameterNames, SyntaxNode bodyNode, String generatedClassName) {
        this.functionName = functionName;
        this.parameterNames = new ArrayList<>(parameterNames);
        this.bodyNode = bodyNode;
        this.generatedClassName = generatedClassName;
    }

    ByteCodeAppender asAppender() {
        return new Appender(functionName, parameterNames, bodyNode, generatedClassName);
    }

    private static final class Appender implements ByteCodeAppender {
        private final String name;
        private final List<String> paramNames;
        private final SyntaxNode bodyNode;
        private final String generatedClassName;
        private final ExpressionEmitter exprEmitter = new BasicExpressionEmitter();
        private final StatementEmitter stmtEmitter = new BasicStatementEmitter(exprEmitter);

        Appender(String name, List<String> paramNames, SyntaxNode bodyNode, String generatedClassName) {
            this.name = name;
            this.paramNames = paramNames;
            this.bodyNode = bodyNode;
            this.generatedClassName = generatedClassName;
        }

        @Override
        public Size apply(MethodVisitor mv, Implementation.Context implementationContext, MethodDescription instrumentedMethod) {
            mv.visitCode();

            LocalVariableAllocator locals = new LocalVariableAllocator(paramNames);

            StackManipulation body = stmtEmitter.emitBody(name, bodyNode, generatedClassName, locals, exprEmitter);
            // Ensure body is not Trivial - if it is, use default return
            if (body == StackManipulation.Trivial.INSTANCE) {
                body = new StackManipulation.Compound(BytecodeOps.pushInt(0), BytecodeOps.ireturn());
            }
            StackManipulation.Size size = body.apply(mv, implementationContext);

            // Use -1, -1 to let ASM automatically compute maxStack and maxLocals
            // This may also help with automatic frame computation
            mv.visitMaxs(-1, -1);
            
            // maxStack calculation: ensure we have enough stack space for our operations
            // Typical operations need: 2 for binary ops (left + right), 1 for comparisons,
            // up to 2-3 for function calls with args. Using 6 as a safe default ensures
            // we have enough space for operations like add(x, y) which needs 2 args + result.
            // Also accounts for nested expressions and temporary stack values during evaluation.
            int maxStack = Math.max(size.getMaximalSize(), 6);
            int maxLocals = locals.getMaxLocals(paramNames.size());
            return new Size(maxStack, maxLocals);
        }
    }
}


